package com.github.cossbow.nsq;

import com.github.cossbow.nsq.util.LambdaUtil;
import com.github.cossbow.nsq.util.ThrowoutConsumer;
import com.github.cossbow.nsq.exceptions.BadMessageException;
import com.github.cossbow.nsq.exceptions.BadTopicException;
import com.github.cossbow.nsq.exceptions.NSQException;
import com.github.cossbow.nsq.frames.ErrorFrame;
import com.github.cossbow.nsq.frames.NSQFrame;
import io.netty.buffer.ByteBufOutputStream;
import reactor.core.publisher.Mono;
import reactor.pool.InstrumentedPool;
import reactor.pool.PoolBuilder;

import java.io.IOException;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Stream;

import static com.github.cossbow.nsq.NSQCommand.multiPublish;
import static com.github.cossbow.nsq.NSQCommand.publish;

public class NSQProducer {
    private final Set<ServerAddress> addresses = ConcurrentHashMap.newKeySet();
    private final AtomicInteger roundRobinCount = new AtomicInteger(0);
    private volatile boolean started = false;
    private InstrumentedPool<Connection<?>> connPool;
    private NSQConfig config = new NSQConfig();

    public synchronized NSQProducer start() {
        if (!started) {
            started = true;
            createPool();
        }
        return this;
    }

    private final ServerAddress[] emptyServerAddresses = new ServerAddress[0];

    private void createPool() {
        Mono<Connection<?>> creator = Mono.fromFuture(() -> {
            var arr = addresses.toArray(emptyServerAddresses);
            if (0 == arr.length) {
                return CompletableFuture.failedFuture(new IllegalStateException("no address"));
            }
            var i = roundRobinCount.getAndIncrement() % arr.length;
            var addr = arr[i];
            return Connection.connect(addr, config,
                    LambdaUtil.doNothing(),
                    LambdaUtil.doNothing());
        });
        connPool = PoolBuilder.from(creator)
                .destroyHandler(c -> Mono.fromCompletionStage(c.closeAsync()))
                .evictionPredicate((c, md) -> !c.isHealthy())
                .buildPool();
    }

    public Map<String, Number> metrics() {
        var m = connPool.metrics();
        return Map.of(
                "acquired", m.acquiredSize(),
                "allocated", m.allocatedSize(),
                "idleSize", m.idleSize(),
                "maxAllocated", m.getMaxAllocatedSize()
        );
    }

    <R> CompletableFuture<R> withConnection(
            Function<Connection<?>, CompletableFuture<R>> callback) {
        var result = new CompletableFuture<R>();
        connPool.acquire().retry(5).doOnCancel(() -> {
            result.completeExceptionally(new CancellationException());
        }).doOnError(result::completeExceptionally).doOnSuccess(ref -> {
            if (null == ref) {
                result.completeExceptionally(
                        new IllegalStateException("borrow null ref"));
                return;
            }
            try {
                var connection = ref.poolable();
                var re = callback.apply(connection);
                re.whenComplete((t, e) -> {
                    ref.release().subscribe();
                    if (null == e) {
                        result.complete(t);
                    } else {
                        result.completeExceptionally(e);
                    }
                });
            } catch (Throwable e) {
                ref.release().subscribe();
                result.completeExceptionally(e);
            }
        }).subscribe();
        return result;
    }

    public CompletableFuture<NSQFrame> send(NSQCommand command) {
        return withConnection(connection ->
                connection.sendAndRecv(command));
    }

    /**
     * produce multiple messages.
     */
    public void produceMulti(String topic, List<byte[]> messages) throws NSQException {
        assertStarted();
        if (messages == null || messages.isEmpty()) {
            return;
        }

        if (messages.size() == 1) {
            //encoding will be screwed up if we MPUB a
            this.produce(topic, messages.get(0));
            return;
        }

        var future = send(multiPublish(topic, messages));
        var frame = future.join();
        checkErrorFrame(frame);
    }

    public void produceMulti(
            String topic,
            Stream<ThrowoutConsumer<ByteBufOutputStream, IOException>> callbackStream)
            throws NSQException {
        assertStarted();
        if (callbackStream == null) return;

        var future = send(multiPublish(topic, callbackStream));
        var frame = future.join();
        checkErrorFrame(frame);
    }

    public void produce(String topic, byte[] message) throws NSQException {
        assertStarted();

        var future = send(publish(topic, message));
        var frame = future.join();
        checkErrorFrame(frame);
    }

    public CompletableFuture<Void> produceAsync(String topic, int defer, byte[] message) {
        assertStarted();
        return withConnection(connection -> {
            var command = defer > 0 ? publish(topic, defer, message) : publish(topic, message);
            return connection.send(command);
        });
    }

    public CompletableFuture<Void> produceAsync(String topic, byte[] message) {
        return produceAsync(topic, 0, message);
    }

    public CompletableFuture<Void> produceAsync(String topic, int defer, ThrowoutConsumer<ByteBufOutputStream, IOException> callback) {
        assertStarted();

        return withConnection(connection -> {
            var command = defer > 0 ? publish(topic, defer, callback) : publish(topic, callback);
            return connection.send(command);
        });
    }

    public CompletableFuture<Void> produceAsync(String topic, ThrowoutConsumer<ByteBufOutputStream, IOException> callback) {
        return produceAsync(topic, 0, callback);
    }

    private void assertStarted() {
        if (started) return;
        throw new IllegalStateException("Producer must be started before producing messages!");
    }

    //

    public NSQProducer addAddress(String host, int port) {
        addresses.add(new ServerAddress(host, port));
        return this;
    }

    public NSQProducer addAddress(ServerAddress address) {
        addresses.add(address);
        return this;
    }

    public NSQProducer addAddresses(Collection<ServerAddress> addresses) {
        this.addresses.addAll(addresses);
        return this;
    }

    public NSQProducer removeAddress(String host, int port) {
        addresses.remove(new ServerAddress(host, port));
        return this;
    }


    private void checkErrorFrame(NSQFrame frame) throws BadMessageException, BadTopicException {
        if (frame instanceof ErrorFrame) {
            String err = frame.readData();
            if (err.startsWith("E_BAD_TOPIC")) {
                throw new BadTopicException(err);
            }
            if (err.startsWith("E_BAD_MESSAGE")) {
                throw new BadMessageException(err);
            }
        }
    }

    public NSQProducer setConfig(NSQConfig config) {
        if (!started) {
            this.config = config;
        }
        return this;
    }


    public void shutdown() {
        started = false;
        connPool.dispose();
    }
}
